package com.star.springai.chat.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AiConfig {

    @Autowired
    private ChatModel chatModel;
    @Autowired
    private VectorStore vectorStore;
    @Autowired
    JdbcChatMemoryRepository chatMemoryRepository;


    @Bean
    public ChatClient chatClient() {

        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(chatMemory())
                .build();
        return ChatClient.builder(chatModel)
                .defaultAdvisors(chatMemoryAdvisor) // 移除questionAnswerAdvisor，让AI更自由发挥
                .defaultSystem(a -> a.text("你是一个智能助手，能够回答各种问题并提供帮助。你可以自由发挥，给出有创意和有用的回答。"))
                .build();

    }


    @Bean
    public QuestionAnswerAdvisor questionAnswerAdvisor() {
        return QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(SearchRequest.builder()
                        .topK(3) // 只返回最相关的3个结果
                        .similarityThreshold(0.7) // 设置相似度阈值，避免不相关的结果
                        .build())
                .build();
    }

    @Bean
    public Advisor retrievalAugmentationAdvisor() {
        return RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.9)
                        .topK(3)
                        .vectorStore(vectorStore)
                        .build())
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(true)
                        .build())
                .build();
    }


    
    // 创建一个带向量存储增强的ChatClient（可选使用）
    @Bean("enhancedChatClient")
    public ChatClient getEnhancedChatClient() {

        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(chatMemory())
                .build();
        return ChatClient.builder(chatModel)
                .defaultAdvisors(retrievalAugmentationAdvisor(), chatMemoryAdvisor)
                .defaultSystem(a -> a.text("你是一个智能助手，能够结合知识库回答问题。优先使用知识库（上下文）中的信息，如果上下文中没有相关信息，请搜索或者根据你自己的认知回答。"))
                .build();
    }

    @Bean
    public ChatMemory chatMemory() {
        return MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .maxMessages(20) // 增加记忆容量，让AI有更多上下文
                .build();
    }
}
